import streamlit as st
import requests
import json
import time
from typing import List, Dict, Any
from streamlit_ace import st_ace
import re

# 页面配置
st.set_page_config(
    page_title="Cairo 智能合约助手",
    page_icon="🔷",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 禁用Streamlit分析功能，避免Segment.io网络错误
import os
os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"

# 自定义CSS样式
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(135deg, #1890ff 0%, #722ed1 100%);
        padding: 2rem;
        border-radius: 10px;
        color: white;
        text-align: center;
        margin-bottom: 2rem;
    }
    
    .chat-message {
        padding: 1rem;
        margin: 0.5rem 0;
        border-radius: 10px;
        border-left: 4px solid #1890ff;
    }
    
    .user-message {
        background-color: #f0f8ff;
        border-left-color: #1890ff;
    }
    
    .assistant-message {
        background-color: #f8f9fa;
        border-left-color: #52c41a;
    }
    
    .code-block {
        background-color: #2d3748;
        color: #e2e8f0;
        padding: 1rem;
        border-radius: 8px;
        margin: 0.5rem 0;
        font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
        overflow-x: auto;
    }
    
    .deep-research {
        background-color: #fff7e6;
        border: 1px solid #ffd591;
        padding: 0.5rem;
        border-radius: 6px;
        margin: 0.5rem 0;
    }
    
    .status-indicator {
        display: flex;
        align-items: center;
        gap: 0.5rem;
        color: #666;
        font-style: italic;
    }
    
    .external-source {
        background-color: #e8f4fd;
        border-left: 4px solid #1890ff;
        padding: 0.5rem;
        margin: 0.5rem 0;
        border-radius: 4px;
    }
    
    .source-badge {
        display: inline-block;
        background-color: #1890ff;
        color: white;
        padding: 0.2rem 0.5rem;
        border-radius: 12px;
        font-size: 0.8rem;
        margin-right: 0.5rem;
    }
    
    .github-badge { background-color: #24292e; }
    .stackoverflow-badge { background-color: #f48024; }
    .docs-badge { background-color: #52c41a; }
    
    .search-progress {
        background-color: #f0f8ff;
        border: 1px solid #d1ecf1;
        padding: 1rem;
        border-radius: 8px;
        margin: 1rem 0;
    }
    
    .tool-status {
        display: flex;
        align-items: center;
        gap: 0.5rem;
        margin: 0.3rem 0;
    }
    
    .tool-status.searching {
        color: #1890ff;
    }
    
    .tool-status.completed {
        color: #52c41a;
    }
    
    .tool-status.error {
        color: #ff4d4f;
    }
    
    .streaming-text {
        animation: typing 0.5s ease-in-out;
    }
    
    @keyframes typing {
        from { opacity: 0.7; }
        to { opacity: 1; }
    }
    
    .streaming-cursor {
        display: inline-block;
        background-color: #1890ff;
        width: 2px;
        height: 1.2em;
        animation: blink 1s infinite;
        margin-left: 2px;
    }
    
    @keyframes blink {
        0%, 50% { opacity: 1; }
        51%, 100% { opacity: 0; }
    }
</style>
""", unsafe_allow_html=True)

# 初始化session state
if "messages" not in st.session_state:
    st.session_state.messages = []
if "is_loading" not in st.session_state:
    st.session_state.is_loading = False

# 后端API配置
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

try:
    from config.settings import config
    API_BASE_URL = f"http://localhost:{config.api.port}"
except ImportError:
    # 如果导入失败，使用默认配置
    API_BASE_URL = "http://localhost:8000"

def call_chat_api(message: str, deep_research: bool = False):
    """调用后端聊天API，返回生成器"""
    try:
        response = requests.post(
            f"{API_BASE_URL}/api/chat/stream",
            json={
                "message": message,
                "deep_research": deep_research
            },
            stream=True,
            timeout=30
        )
        
        if response.status_code != 200:
            yield {"type": "error", "content": f"❌ API错误: {response.status_code}"}
            return
        
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')
                if line_str.startswith('data: '):
                    try:
                        data = json.loads(line_str[6:])
                        yield data
                    except json.JSONDecodeError:
                        continue
        
    except requests.exceptions.RequestException as e:
        yield {"type": "error", "content": f"❌ 网络错误: {str(e)}"}
    except Exception as e:
        yield {"type": "error", "content": f"❌ 未知错误: {str(e)}"}

def format_code_blocks(text: str) -> str:
    """格式化代码块"""
    # 检测代码块
    code_pattern = r'```(\w+)?\n?(.*?)```'
    
    def replace_code(match):
        language = match.group(1) or 'rust'
        code = match.group(2).strip()
        return f'<div class="code-block"><pre><code class="language-{language}">{code}</code></pre></div>'
    
    return re.sub(code_pattern, replace_code, text, flags=re.DOTALL)

def display_message(message: Dict[str, Any], is_streaming: bool = False):
    """显示消息（支持外部工具结果和流式显示）"""
    role = message.get("role", "user")
    content = message.get("content", "")
    deep_research = message.get("deep_research", False)
    timestamp = message.get("timestamp", time.time())
    metadata = message.get("metadata", {})
    
    # 格式化时间
    time_str = time.strftime("%H:%M:%S", time.localtime(timestamp))
    
    with st.container():
        if role == "user":
            st.markdown(f"""
            <div class="chat-message user-message">
                <strong>👤 用户</strong> <span style="color: #999; font-size: 0.8rem;">{time_str}</span>
                <div style="margin-top: 0.5rem;">{content}</div>
            </div>
            """, unsafe_allow_html=True)
        else:
            research_badge = ""
            if deep_research:
                research_badge = '<span class="deep-research">🔬 深度研究</span>'
            
            # 添加外部工具标识
            external_tools_badge = ""
            if metadata.get("has_external_results"):
                external_tools_badge = '<span class="source-badge">🔍 外部搜索</span>'
            
            # 处理代码块高亮
            formatted_content = format_code_blocks(content)
            
            # 增强外部链接显示
            formatted_content = enhance_external_links(formatted_content)
            
            # 为流式消息添加特殊样式和光标
            content_class = "streaming-text" if is_streaming else ""
            cursor = '<span class="streaming-cursor"></span>' if is_streaming else ""
            
            st.markdown(f"""
            <div class="chat-message assistant-message">
                <strong>🤖 助手</strong> <span style="color: #999; font-size: 0.8rem;">{time_str}</span> {research_badge} {external_tools_badge}
                <div class="{content_class}" style="margin-top: 0.5rem;">{formatted_content}{cursor}</div>
            </div>
            """, unsafe_allow_html=True)


def enhance_external_links(content: str) -> str:
    """增强外部链接显示"""
    # 为不同来源的链接添加样式
    content = re.sub(
        r'\*\*\[Cairo 官方文档\]\*\*',
        '<span class="source-badge docs-badge">📚 Cairo 官方文档</span>',
        content
    )
    content = re.sub(
        r'\*\*\[Starknet 官方文档\]\*\*',
        '<span class="source-badge docs-badge">📖 Starknet 官方文档</span>',
        content
    )
    content = re.sub(
        r'\*\*\[GitHub 代码\]\*\*',
        '<span class="source-badge github-badge">💻 GitHub 代码</span>',
        content
    )
    content = re.sub(
        r'\*\*\[GitHub Issues\]\*\*',
        '<span class="source-badge github-badge">🐛 GitHub Issues</span>',
        content
    )
    content = re.sub(
        r'\*\*\[Stack Overflow\]\*\*',
        '<span class="source-badge stackoverflow-badge">❓ Stack Overflow</span>',
        content
    )
    
    return content


def display_search_progress(tools_status: Dict[str, str]):
    """显示搜索进度"""
    st.markdown("""
    <div class="search-progress">
        <h4>🔍 正在搜索外部资源...</h4>
    </div>
    """, unsafe_allow_html=True)
    
    for tool, status in tools_status.items():
        tool_names = {
            "cairo_docs": "📚 Cairo 官方文档",
            "github": "💻 GitHub 代码库", 
            "stackoverflow": "❓ Stack Overflow"
        }
        
        tool_display = tool_names.get(tool, tool)
        status_class = f"tool-status {status}"
        
        if status == "searching":
            icon = "🔄"
            text = "搜索中..."
        elif status == "completed":
            icon = "✅"
            text = "完成"
        else:
            icon = "❌"
            text = "错误"
        
        st.markdown(f"""
        <div class="{status_class}">
            {icon} {tool_display}: {text}
        </div>
        """, unsafe_allow_html=True)

# 主界面
def main():
    # 标题
    st.markdown("""
    <div class="main-header">
        <h1>🔷 Cairo 智能合约助手</h1>
        <p>基于 LangGraph 的垂直领域聊天机器人 | 千问3驱动</p>
    </div>
    """, unsafe_allow_html=True)
    
    # 侧边栏
    with st.sidebar:
        st.header("⚙️ 设置")
        
        # 模型信息
        st.info("""
        **当前配置:**
        - 模型: 千问3 (qwen-plus)
        - 模式: 流式对话
        - 支持: 深度研究
        """)
        
        # API连接测试
        if st.button("🔗 测试API连接"):
            try:
                response = requests.get(f"{API_BASE_URL}/api/health", timeout=5)
                if response.status_code == 200:
                    st.success("✅ API连接正常")
                else:
                    st.error(f"❌ API错误: {response.status_code}")
            except Exception as e:
                st.error(f"❌ 连接失败: {str(e)}")
        
        # 功能说明
        st.markdown("""
        **功能特性:**
        - ✅ Cairo语法解析
        - ✅ 代码错误诊断
        - ✅ 最佳实践建议
        - ✅ 深度研究分析
        - ✅ 代码高亮显示
        """)
        
        # 清空对话
        if st.button("🗑️ 清空对话", type="secondary"):
            st.session_state.messages = []
            st.rerun()
    
    # 聊天区域
    col1, col2 = st.columns([4, 1])
    
    with col1:
        # 输入框
        user_input = st.text_area(
            "💬 请输入您关于Cairo的问题...",
            height=100,
            placeholder="例如: Cairo中的felt252类型有什么特点？",
            disabled=st.session_state.is_loading
        )
    
    with col2:
        st.write("")  # 空行
        st.write("")  # 空行
        
        # 发送按钮
        send_button = st.button("📤 发送问题", type="primary", disabled=st.session_state.is_loading)
        
        # 深度研究按钮
        deep_research_button = st.button(
            "🔍 深度研究", 
            type="secondary", 
            disabled=st.session_state.is_loading,
            help="获得更详细的分析报告"
        )
    
    # 处理用户输入
    if (send_button or deep_research_button) and user_input.strip():
        # 添加用户消息
        st.session_state.messages.append({
            "role": "user",
            "content": user_input,
            "timestamp": time.time()
        })
        
        # 设置加载状态
        st.session_state.is_loading = True
        
        # 调用API
        is_deep_research = deep_research_button
        try:
            # 显示当前用户消息
            display_message(st.session_state.messages[-1])
            
            # 创建独立的容器用于流式显示
            status_container = st.empty()
            response_container = st.empty()
            
            # 流式显示响应
            assistant_response = ""
            chunk_count = 0
            external_search_status = {}
            has_external_results = False
            
            # 显示初始状态
            status_container.markdown('<div class="status-indicator">⏳ 正在连接API...</div>', unsafe_allow_html=True)
            
            # 调用API并处理响应
            for data in call_chat_api(user_input, is_deep_research):
                event_type = data.get('type', '')
                content = data.get('content', '')
                
                # 处理不同类型的事件
                if event_type == 'status':
                    status_container.markdown(f'<div class="status-indicator">⏳ {content}</div>', unsafe_allow_html=True)
                
                elif event_type == 'external_search':
                    # 显示外部搜索状态
                    sources = data.get('sources', [])
                    search_query = data.get('query', '相关内容')
                    search_stage = data.get('stage', 'searching')
                    has_external_results = True
                    
                    # 根据搜索阶段显示不同的状态
                    if search_stage == 'starting':
                        search_status_html = f'<div class="search-progress"><h4>🔍 开始搜索外部资源...</h4><p>搜索关键词: {search_query}</p>'
                    elif search_stage == 'searching':
                        search_status_html = f'<div class="search-progress"><h4>🔍 正在搜索外部资源...</h4><p>搜索关键词: {search_query}</p>'
                    elif search_stage == 'completed':
                        search_status_html = f'<div class="search-progress"><h4>✅ 外部搜索完成</h4><p>搜索关键词: {search_query}</p>'
                    else:
                        search_status_html = f'<div class="search-progress"><h4>🔍 正在搜索外部资源...</h4><p>搜索关键词: {search_query}</p>'
                    
                    # 显示各个搜索源的状态
                    for source in sources:
                        source_status = data.get(f'{source}_status', 'searching')
                        if source == 'github':
                            if source_status == 'completed':
                                search_status_html += '<div class="tool-status completed">✅ GitHub 代码搜索完成</div>'
                            elif source_status == 'error':
                                search_status_html += '<div class="tool-status error">❌ GitHub 搜索失败 (未配置Token)</div>'
                            else:
                                search_status_html += '<div class="tool-status searching">🔄 正在搜索 GitHub 代码...</div>'
                        elif source == 'stackoverflow':
                            if source_status == 'completed':
                                search_status_html += '<div class="tool-status completed">✅ Stack Overflow 搜索完成</div>'
                            elif source_status == 'error':
                                search_status_html += '<div class="tool-status error">❌ Stack Overflow 搜索失败</div>'
                            else:
                                search_status_html += '<div class="tool-status searching">🔄 正在搜索 Stack Overflow...</div>'
                        elif source == 'cairo_docs':
                            if source_status == 'completed':
                                search_status_html += '<div class="tool-status completed">✅ Cairo 文档搜索完成</div>'
                            elif source_status == 'timeout':
                                search_status_html += '<div class="tool-status error">⏰ Cairo 文档搜索超时</div>'
                            elif source_status == 'error':
                                search_status_html += '<div class="tool-status error">❌ Cairo 文档搜索失败</div>'
                            else:
                                search_status_html += '<div class="tool-status searching">🔄 正在搜索 Cairo 文档...</div>'
                    search_status_html += '</div>'
                    status_container.markdown(search_status_html, unsafe_allow_html=True)
                
                elif event_type in ['chunk', 'structured_response']:
                    chunk_count += 1
                    assistant_response += content
                    
                    # 实时更新流式内容显示
                    if assistant_response:
                        temp_message = {
                            "role": "assistant",
                            "content": assistant_response,
                            "deep_research": is_deep_research,
                            "metadata": {"has_external_results": has_external_results}
                        }
                        
                        # 使用独立容器显示正在生成的消息，添加流式效果
                        with response_container.container():
                            display_message(temp_message, is_streaming=True)
                        
                        # 更新状态指示器
                        if event_type == 'chunk':
                            status_container.markdown(f'<div class="status-indicator">⏳ 正在生成回答... (已接收 {chunk_count} 个片段)</div>', unsafe_allow_html=True)
                        else:
                            status_container.empty()
                
                elif event_type == 'done':
                    status_container.empty()
                    break
                
                elif event_type == 'error':
                    status_container.empty()
                    response_container.error(f"❌ {content}")
                    break
            
            # 检查是否有响应
            if not assistant_response.strip():
                st.error("❌ 没有收到任何响应，请检查后端服务是否正常运行")
                return
            
            # 添加完整的助手消息
            st.session_state.messages.append({
                "role": "assistant",
                "content": assistant_response,
                "deep_research": is_deep_research,
                "timestamp": time.time(),
                "metadata": {"has_external_results": has_external_results}
            })
            
        except Exception as e:
            st.error(f"❌ 发送消息失败: {str(e)}")
            st.error(f"错误详情: {type(e).__name__}")
        finally:
            st.session_state.is_loading = False
            st.rerun()
    
    # 显示历史消息
    if st.session_state.messages:
        st.markdown("---")
        st.subheader("💬 对话历史")
        
        for message in st.session_state.messages:
            display_message(message)
    else:
        # 欢迎信息
        st.markdown("""
        <div style="text-align: center; padding: 2rem; color: #666;">
            <h3>👋 欢迎使用 Cairo 智能合约助手！</h3>
            <p>我可以帮助您：</p>
            <ul style="text-align: left; display: inline-block;">
                <li>解答 Cairo 语言语法和特性问题</li>
                <li>分析和解释 Cairo 代码片段</li>
                <li>诊断代码错误并提供解决方案</li>
                <li>提供最佳实践建议</li>
            </ul>
            <p><strong>💡 点击"深度研究"按钮可获得更详细的分析报告</strong></p>
        </div>
        """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()
